"""
SSVEP Wang dataset.
"""

import logging
from os.path import dirname

import numpy as np
from mne import create_info
from mne.channels import make_standard_montage
from mne.io import RawArray
from pyunpack import Archive
from scipy.io import loadmat

from . import download as dl
from .base import BaseDataset


log = logging.getLogger()

# WANG_URL = 'http://bci.med.tsinghua.edu.cn/upload/yijun/' # 403 error
# WANG_URL = 'ftp://anonymous@sccn.ucsd.edu/pub/ssvep_benchmark_dataset/'
WANG_URL = "http://www.thubci.com/uploads/down/"


class Wang2016(BaseDataset):
    """SSVEP Wang 2016 dataset

    This dataset gathered SSVEP-BCI recordings of 35 healthy subjects (17
    females, aged 17-34 years, mean age: 22 years) focusing on 40 characters
    flickering at different frequencies (8-15.8 Hz with an interval of 0.2 Hz).
    For each subject, the experiment consisted of 6 blocks. Each block
    contained 40 trials corresponding to all 40 characters indicated in a
    random order. Each trial started with a visual cue (a red square)
    indicating a target stimulus. The cue appeared for 0.5 s on the screen.
    Subjects were asked to shift their gaze to the target as soon as possible
    within the cue duration. Following the cue offset, all stimuli started to
    flicker on the screen concurrently and lasted 5 s. After stimulus offset,
    the screen was blank for 0.5 s before the next trial began, which allowed
    the subjects to have short breaks between consecutive trials. Each trial
    lasted a total of 6 s. To facilitate visual fixation, a red triangle
    appeared below the flickering target during the stimulation period. In each
    block, subjects were asked to avoid eye blinks during the stimulation
    period. To avoid visual fatigue, there was a rest for several minutes
    between two consecutive blocks.

    EEG data were acquired using a Synamps2 system (Neuroscan, Inc.) with a
    sampling rate of 1000 Hz. The amplifier frequency passband ranged from
    0.15 Hz to 200 Hz. Sixty-four channels covered the whole scalp of the
    subject and were aligned according to the international 10-20 system.
    The ground was placed on midway between Fz and FPz. The reference was
    located on the vertex. Electrode impedances were kept below 10 KΩ. To
    remove the common power-line noise, a notch filter at 50 Hz was applied
    in data recording. Event triggers generated by the computer to the
    amplifier and recorded on an event channel synchronized to the EEG data.

    The continuous EEG data was segmented into 6 s epochs (500 ms pre-stimulus,
    5.5 s post-stimulus onset). The epochs were subsequently downsampled to
    250 Hz. Thus each trial consisted of 1500 time points. Finally, these data
    were stored as double-precision floating-point values in MATLAB and were
    named as subject indices (i.e., S01.mat, …, S35.mat). For each file, the
    data loaded in MATLAB generate a 4-D matrix named ‘data’ with dimensions
    of [64, 1500, 40, 6]. The four dimensions indicate ‘Electrode index’,
    ‘Time points’, ‘Target index’, and ‘Block index’. The electrode positions
    were saved in a ‘64-channels.loc’ file. Six trials were available for each
    SSVEP frequency. Frequency and phase values for the 40 target indices were
    saved in a ‘Freq_Phase.mat’ file.

    Information for all subjects was listed in a ‘Sub_info.txt’ file. For each
    subject, there are five factors including ‘Subject Index’, ‘Gender’, ‘Age’,
    ‘Handedness’, and ‘Group’. Subjects were divided into an ‘experienced’
    group (eight subjects, S01-S08) and a ‘naive’ group (27 subjects, S09-S35)
    according to their experience in SSVEP-based BCIs.

    references
    ----------
    .. [1] Y. Wang, X. Chen, X. Gao and S. Gao, "A Benchmark Dataset for
    SSVEP-Based Brain–Computer Interfaces," in IEEE Transactions on Neural
    Systems and Rehabilitation Engineering, vol. 25, no. 10, pp. 1746-1752,
    Oct. 2017
    doi: 10.1109/TNSRE.2016.2627556.
    """

    def __init__(self):
        super().__init__(
            subjects=list(range(1, 35)),
            sessions_per_subject=1,
            # fmt: off
            events={'8': 1, '9': 2, '10': 3, '11': 4, '12': 5, '13': 6,
                    '14': 7, '15': 8, '8.2': 9, '9.2': 10, '10.2': 11,
                    '11.2': 12, '12.2': 13, '13.2': 14, '14.2': 15, '15.2': 16,
                    '8.4': 17, '9.4': 18, '10.4': 19, '11.4': 20, '12.4': 21,
                    '13.4': 22, '14.4': 23, '15.4': 24, '8.6': 25, '9.6': 26,
                    '10.6': 27, '11.6': 28, '12.6': 29, '13.6': 30,
                    '14.6': 31, '15.6': 32, '8.8': 33, '9.8': 34, '10.8': 35,
                    '11.8': 36, '12.8': 37, '13.8': 38, '14.8': 39,
                    '15.8': 40},
            # fmt: on
            code="SSVEP Wang",
            interval=[0.5, 5.5],
            paradigm="ssvep",
            doi="doi://10.1109/TNSRE.2016.2627556",
        )

    def _get_single_subject_data(self, subject):
        """Return the data of a single subject"""
        n_samples, n_channels, n_trials = 1500, 64, 6
        n_classes = len(self.event_id)

        fname = self.data_path(subject)
        Archive(fname).extractall(dirname(fname))
        mat = loadmat(fname[:-4])

        data = np.transpose(mat["data"], axes=(2, 3, 0, 1))
        data = np.reshape(data, newshape=(-1, n_channels, n_samples))
        data = data - data.mean(axis=2, keepdims=True)
        raw_events = np.zeros((data.shape[0], 1, n_samples))
        raw_events[:, 0, 0] = np.array(
            [n_trials * [i + 1] for i in range(n_classes)]
        ).flatten()
        data = np.concatenate([1e-6 * data, raw_events], axis=1)
        # add buffer in between trials
        log.warning(
            "Trial data de-meaned and concatenated with a buffer"
            " to create continuous data"
        )
        buff = (data.shape[0], n_channels + 1, 50)
        data = np.concatenate([np.zeros(buff), data, np.zeros(buff)], axis=2)

        # fmt: off
        ch_names = ['Fp1', 'Fpz', 'Fp2', 'AF3', 'AF4', 'F7', 'F5', 'F3', 'F1',
                    'Fz', 'F2', 'F4', 'F6', 'F8', 'FT7', 'FC5', 'FC3', 'FC1',
                    'FCz', 'FC2', 'FC4', 'FC6', 'FT8', 'T7', 'C5', 'C3', 'C1',
                    'Cz', 'C2', 'C4', 'C6', 'T8', 'M1', 'TP7', 'CP5', 'CP3',
                    'CP1', 'CPz', 'CP2', 'CP4', 'CP6', 'TP8', 'M2', 'P7', 'P5',
                    'P3', 'P1', 'Pz', 'P2', 'P4', 'P6', 'P8', 'PO7', 'PO5',
                    'PO3', 'POz', 'PO4', 'PO6', 'PO8', 'CB1', 'O1', 'Oz', 'O2',
                    'CB2', 'stim']
        # fmt: on
        ch_types = ["eeg"] * 59 + ["misc"] + 3 * ["eeg"] + ["misc", "stim"]
        sfreq = 250
        info = create_info(ch_names, sfreq, ch_types)
        raw = RawArray(data=np.concatenate(list(data), axis=1), info=info, verbose=False)
        montage = make_standard_montage("standard_1005")
        raw.set_montage(montage)
        return {"session_0": {"run_0": raw}}

    def data_path(
        self, subject, path=None, force_update=False, update_path=None, verbose=None
    ):
        if subject not in self.subject_list:
            raise (ValueError("Invalid subject number"))
        url = "{:s}s{:d}.rar".format(WANG_URL, subject)
        return dl.data_path(url, "WANG", path, force_update, update_path, verbose)
